module mod_chem_reactions

   use iso_fortran_env, only: dp => real64
   use stdlib_strings
   use stdlib_string_type
   use stdlib_str2num
   use stdlib_optval
   use mod_strings_ext
   use mod_chem_species, only: type_chem_specy
   use mod_chem_unit

   private
   public :: type_chem_reaction

#include <chem.f90>
   type :: type_chem_reaction

      integer :: n_species
      character(len=:), allocatable :: equation
      character(len=:), allocatable :: type
      real(dp) :: A
      real(dp) :: b
      real(dp) :: Ea
      real(dp) :: A_unit_factor = -1.0_dp
      real(dp) :: Ea_unit_factor = -1.0_dp
      real(dp) :: kf, kb !< fwd/rev rate constant
      real(dp) :: Rf, Rb !< fwd/rev reaction rate
      character(len=:), allocatable :: A_unit
      character(len=:), allocatable :: b_unit
      character(len=:), allocatable :: Ea_unit
      real(dp), allocatable :: efficiencies(:)
      integer, allocatable :: sto_num_fwd(:)
      integer, allocatable :: sto_num_rev(:)
      type(type_chem_specy), pointer :: species(:)
   contains
      procedure :: parse_reaction => ParseReaction

      procedure :: get_equilibrium_constant_x => GetEquilibriumConstantX
      procedure :: get_equilibrium_constant_p => GetEquilibriumConstantP
      procedure :: get_fwdrate_constant => GetFwdRateConstant
      procedure :: get_revrate_constant => GetRevRateConstant
      procedure :: get_fwdrate_progress => GetFwdRateOfProgress
      procedure :: get_revrate_progress => GetRevRateOfProgress
      procedure :: get_netrate_progress => GetNetRateOfProgress
      
      procedure :: get_fwdrate_progress_dx => GetFwdRateOfProgressDX
      procedure :: get_revrate_progress_dx => GetRevRateOfProgressDX
      procedure :: get_netrate_progresses => GetNetRateOfProgresses

      procedure :: get_creation_production_rates => GetCreationProductionRates
      procedure :: get_desctruction_production_rates => GetDestructionProductionRates
      procedure :: get_net_production_rates => GetNetProductionRates
      procedure :: get_creation_production_rates_dx => GetCreationProductionRatesDX
      procedure :: get_desctruction_production_rates_dx => GetDestructionProductionRatesDX
      procedure :: get_net_production_rates_dx => GetNetProductionRatesDX

      procedure, private :: initialize => ReactionsInitialize
      procedure :: Arrhenius
   end type type_chem_reaction

contains

   subroutine ReactionsInitialize(this)
      class(type_chem_reaction), intent(inout) :: this

      this%A_unit_factor = get_A_unit_factor(this, this%A_unit)
      this%Ea_unit_factor = get_Ea_unit_factor(this, this%Ea_unit)

   end subroutine ReactionsInitialize

   function get_A_unit_factor(this, unit) result(factor)
      class(type_chem_reaction), intent(in) :: this
      character(len=*), intent(in) :: unit
      real(dp) :: factor
      integer :: nsize

      type(string_type), allocatable :: strings(:), strings2(:)

      call split(string_type(unit), strings, "/")
      if (size(strings) /= 3 .or. strings(2) == "s") then
         print *, "Error: the unit for Arrhenius A should be like cm^3/mol/s or cm^6/mol^2/s."
         error stop
      end if

      nsize = count(this%sto_num_fwd > 0)
      if (this%type == "three-body") nsize = nsize + 1

      call split(strings(1), strings2, "^")
      if (to_num(char(strings2(2)), 1) /= 3*(nsize - 1)) then
         print *, "Error: the unit for Arrhenius A should be like &
&                  cm^"//string_type(3*(nsize - 1))//"/mol^"//string_type(nsize - 1)//"/s."
      end if
      factor = get_factor_trans_units(char(strings2(1)), "m", "length")**(3*(nsize - 1))

      call split(strings(2), strings2, "^")
      if (size(strings2) > 1) then
         if (to_num(char(strings2(2)), 1) /= (nsize - 1)) then
            print *, "Error: the unit for Arrhenius A should be like &
&            cm^"//string_type(3*(nsize - 1))//"/mol^"//string_type(nsize - 1)//"/s."
         end if
      end if
      factor = factor/(get_factor_trans_units(char(strings2(1)), "kmol", "quantity")**(nsize - 1))

      factor = factor/get_factor_trans_units(char(strings(3)), "s", "time")

   end function get_A_unit_factor

   function get_Ea_unit_factor(this, unit) result(factor)
      class(type_chem_reaction), intent(in) :: this
      character(len=*), intent(in) :: unit
      real(dp) :: factor

      type(string_type), allocatable :: strings(:)

      call split(string_type(unit), strings, "/")

      factor = get_factor_trans_units(char(strings(1)), "J", "energy")/ &
               get_factor_trans_units(char(strings(2)), "kmol", "quantity")

   end function get_Ea_unit_factor

   function Arrhenius(this, T, Tv) result(k)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp) :: k

      real(dp) :: T_

      T_ = sqrt(T*optval(Tv, T))
      k = this%A_unit_factor*this%A*T_**this%b*exp(-this%Ea*this%Ea_unit_factor/(R0*T_))
      ! k = this%Ea * this%Ea_unit_factor/R0

   end function Arrhenius

   function GetFwdRateConstant(this, T, Tv) result(k)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp) :: k

      integer :: i

      k = this%Arrhenius(T, Tv)

   end function GetFwdRateConstant

   function GetRevRateConstant(this, T, Tv) result(k)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp) :: k

      integer :: i

      k = this%Arrhenius(T)/this%get_equilibrium_constant_x(T)

   end function GetRevRateConstant

   function GetFwdRateOfProgress(this, X, T, Tv) result(R)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp) :: R

      integer :: i

      R = this%get_fwdrate_constant(T, Tv)
      do i = 1, this%n_species
         R = R*X(i)**this%sto_num_fwd(i)
      end do
      if (this%type == "three-body") &
         R = R*sum(this%efficiencies*X)

   end function GetFwdRateOfProgress

   function GetFwdRateOfProgressDX(this, R, X) result(rr)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: R
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), allocatable :: rr(:)

      integer :: i

      allocate(rr(this%n_species))
      do i = 1, this%n_species
         rr(i) = this%sto_num_fwd(i)*R/(X(i)+sqrt(epsilon(1.0_dp)))
      enddo
      ! print*, rr

   end function GetFwdRateOfProgressDX

   function GetRevRateOfProgress(this, X, T, Tv) result(R)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp) :: R

      integer :: i

      R = this%get_revrate_constant(T)
      do i = 1, this%n_species
         R = R*X(i)**this%sto_num_rev(i)
      end do
      if (this%type == "three-body") &
         R = R*sum(this%efficiencies*X)

      R = -R

   end function GetRevRateOfProgress

   function GetRevRateOfProgressDX(this, R, X) result(rr)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: R
      real(dp), allocatable :: rr(:)

      integer :: i

      allocate(rr(this%n_species))
      do i = 1, this%n_species
         rr(i) = this%sto_num_rev(i)*R/(X(i)+sqrt(epsilon(1.0_dp)))
      enddo

   end function GetRevRateOfProgressDX

   function GetNetRateOfProgress(this, X, T, Tv) result(R)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp) :: R

      R = this%get_fwdrate_progress(X, T, Tv) + this%get_revrate_progress(X, T)

   end function GetNetRateOfProgress

   subroutine GetNetRateOfProgresses(this, X, T, Tv, R, rr)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp), intent(out) :: R
      real(dp), intent(inout) :: rr(:)

      real(dp) :: fwdrate, revrate

      fwdrate = this%get_fwdrate_progress(X, T, Tv)
      revrate = this%get_revrate_progress(X, T)

      R = fwdrate + revrate
      rr = this%get_fwdrate_progress_dx(fwdrate, X) + this%get_revrate_progress_dx(revrate, X)

   end subroutine GetNetRateOfProgresses


   function GetEquilibriumConstantX(this, T, Tv) result(Kx)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp) :: Kx
      real(dp), parameter :: atm = 101325.0_dp

      Kx = this%get_equilibrium_constant_p(T, Tv)*(atm/(R0*T))**(sum(this%sto_num_rev - this%sto_num_fwd))

   end function GetEquilibriumConstantX

   function GetEquilibriumConstantP(this, T, Tv) result(Kp)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp) :: Kp

      Kp = exp(-sum((this%sto_num_rev - this%sto_num_fwd)*this%species%get_gibbs_RT(T)))

   end function GetEquilibriumConstantP

   function GetCreationProductionRates(this, X, T, Tv) result(R)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp), allocatable :: R(:)

      integer :: i

      R = (this%sto_num_rev - this%sto_num_fwd)*this%get_revrate_progress(X, T, Tv)

   end function GetCreationProductionRates

   subroutine GetCreationProductionRatesDX(this, X, T, Tv, R, rr)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp), intent(inout) :: R(:)
      real(dp), intent(inout) :: rr(:, :)

      real(dp) :: revrate
      integer :: i

      revrate = this%get_revrate_progress(X, T, Tv)

      R = (this%sto_num_rev - this%sto_num_fwd)*revrate
      do i = 1, this%n_species
         rr(i, :) = (this%sto_num_rev(i) - this%sto_num_fwd(i))*this%get_revrate_progress_dx(revrate, X)
      enddo

   end subroutine GetCreationProductionRatesDX


   function GetDestructionProductionRates(this, X, T, Tv) result(R)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp), allocatable :: R(:)

      integer :: i

      R = (this%sto_num_rev - this%sto_num_fwd)*this%get_fwdrate_progress(X, T, Tv)

   end function GetDestructionProductionRates

   subroutine GetDestructionProductionRatesDX(this, X, T, Tv, R, rr)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp), intent(inout) :: R(:)
      real(dp), intent(inout) :: rr(:, :)

      integer :: i

      real(dp) :: fwdrate

      fwdrate = this%get_fwdrate_progress(X, T, Tv)

      R = (this%sto_num_rev - this%sto_num_fwd)*fwdrate
      do i = 1, this%n_species
         rr(i, :) = (this%sto_num_rev(i) - this%sto_num_fwd(i))*this%get_fwdrate_progress_dx(fwdrate, X)
      enddo

   end subroutine GetDestructionProductionRatesDX

   function GetNetProductionRates(this, X, T, Tv) result(R)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp), allocatable :: R(:)

      integer :: i
      R = this%get_desctruction_production_rates(X, T, Tv) + &
          this%get_creation_production_rates(X, T, Tv)

   end function GetNetProductionRates

   subroutine GetNetProductionRatesDX(this, X, T, Tv, R, rr)
      class(type_chem_reaction), intent(in) :: this
      real(dp), intent(in) :: X(:) !< molar_density kmol/m^3
      real(dp), intent(in) :: T
      real(dp), intent(in), optional :: Tv
      real(dp), intent(inout) :: R(:)
      real(dp), intent(inout) :: rr(:, :)

      integer :: i
      real(dp) :: fwdrate, revrate

      fwdrate = this%get_fwdrate_progress(X, T, Tv)
      revrate = this%get_revrate_progress(X, T, Tv)
      
      R = (this%sto_num_rev - this%sto_num_fwd)*(fwdrate + revrate)
      do i = 1, this%n_species
         rr(i, :) = (this%sto_num_rev(i) - this%sto_num_fwd(i))*(this%get_fwdrate_progress_dx(fwdrate, X) + this%get_revrate_progress_dx(revrate, X))
      enddo

   end subroutine GetNetProductionRatesDX

   subroutine ParseReaction(this, species)
      class(type_chem_reaction), intent(inout) :: this
      type(type_chem_specy), intent(in), target :: species(:)

      integer :: idx
      type(string_type), allocatable :: strings(:)

      call split(string_type(this%equation), strings, " <=> ")
      if (size(strings) == 1) call split(string_type(this%equation), strings, " = ")

      call parse_stoichiometric_number(strings(1), species, this%sto_num_fwd)
      call parse_stoichiometric_number(strings(2), species, this%sto_num_rev)

      this%n_species = size(species)
      this%species => species
      call this%initialize()

   end subroutine ParseReaction

   subroutine parse_stoichiometric_number(string, species, sto_num)

      type(string_type), intent(in) :: string
      type(type_chem_specy), intent(in) :: species(:)
      integer, allocatable, intent(out) :: sto_num(:)

      type(string_type), allocatable :: token(:), strings(:)
      integer :: l, i

      allocate (sto_num(size(species)))
      sto_num = 0

      call split(string, strings, " + ")
      do l = 1, size(strings)
         if (count(strings(l), " ") == 0) then
            do i = 1, size(species)
               if (strings(l) == species(i)%name) then
                  sto_num(i) = 1
               end if
            end do
         else
            call split(strings(l), token, " ")
            do i = 1, size(species)
               if (token(2) == species(i)%name) then
                  sto_num(i) = to_num(char(token(1)), 1)
               end if
            end do
         end if
      end do

   end subroutine parse_stoichiometric_number

end module mod_chem_reactions
